package cn.qiyu5522.tk_ems.interceptor;

import cn.qiyu5522.tk_ems.utils.TraceIdContext;
import cn.qiyu5522.tk_ems.utils.TraceIdUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.web.servlet.HandlerInterceptor;
import org.springframework.web.servlet.ModelAndView;

import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;

/**
 * @BelongsProject: tk_ems
 * @BelongsPackage: cn.qiyu5522.tk_ems.interceptor
 * @Author: TangQi
 * @CreateTime: 2023-02-06  09:43
 * @Description: TODO
 * @Version: 1.0
 */

public class TraceIdRequestLoggingInterceptor implements HandlerInterceptor {
    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        String requestId = request.getHeader(TraceIdContext.TRACE_ID_KEY);
        if(StringUtils.isNotEmpty(requestId)){
            TraceIdContext.setTraceId(requestId);
        }else{
            // 从Header中获取requestId，如果获取不到则视为“源头”，生成一个requestId，设置到MDC当中。
            TraceIdContext.setTraceId(TraceIdUtils.getTraceId());
        }

        return HandlerInterceptor.super.preHandle(request, response, handler);
    }

    @Override
    public void postHandle(HttpServletRequest request, HttpServletResponse response, Object handler, ModelAndView modelAndView) throws Exception {
        HandlerInterceptor.super.postHandle(request, response, handler, modelAndView);
    }

    @Override
    public void afterCompletion(HttpServletRequest request, HttpServletResponse response, Object handler, Exception ex) throws Exception {
        TraceIdContext.removeTraceId();
    }
}
